import time

from numpy import *
from pylab import *

from numpy.random import random
import numpy.linalg as linalg
import scipy.sparse as sparse
import scipy.sparse.csgraph as sgraph
# from scipy.sparse.linalg.dsolve.umfpack import UmfpackContext as umf


def genER(n, p=0.5):
    A = sparse.lil_matrix((n, n), dtype=int8)
    for i in range(n):
        for j in range(i + 1, n):
            if random() < p:
                A[i, j] = A[j, i] = 1
    return A.tocsr()


def genSF(n, m=1, m0=2):
    A = sparse.lil_matrix((n, n), dtype=int8)
    A0 = genER(m0, 1)
    A[:m0, :m0] = A0
    degrees = A.sum(1)
    sumDegree = degrees.sum()
    for k in range(m0, n):
        connected = []
        totalDegree = sumDegree
        for _ in range(m):
            rnd = random() * totalDegree
            curDegree = 0
            for i in range(k):
                if i not in connected:
                    curDegree += degrees[i]
                    if rnd < curDegree:
                        connected.append(i)
                        totalDegree -= degrees[i]
                        degrees[i] += 1
                        A[i, k] = A[k, i] = 1
                        break
        degrees[k] += m
        sumDegree += 2 * m
    return A.tocsr()


#TODO: implement for sparse matrix, using Cholesky or LU decompose
def determinant(A):
    return linalg.det(A)


#TODO: review the approach of obtaining minor
def cofactor(A, i=0, j=0):
    n = A.shape[0]
    return (-1) ** (i + j) * determinant(A[array(list(range(i)) + list(range(i + 1, n)))[:, newaxis],
                                           array(list(range(j)) + list(range(j + 1, n)))])


def cofactor2(A, i=0, j=0):
    n = A.shape[0]
    if i > j:
        i, j = j, i
    chain = array(list(range(i)) + list(range(i + 1, j)) + list(range(j + 1, n)))
    return (-1) ** (i + j) * determinant(A[chain[:, newaxis], chain])


def resistance(T):
    n = T.shape[0]
    const = cofactor(T)
    # print("Kirgoff const:", const)   # TODO: define why const is so huge
    R = sparse.lil_matrix((n, n))
    for i in range(n):
        for j in range(i, n):
            if i == j:
                R[i, i] = 1
            else:
                R[i, j] = R[j, i] = cofactor2(T, i, j) / const
    return R.tocsr()


def hiddenLinks(A, R, threshold=0.5):
    n = A.shape[0]
    iA = sparse.csr_matrix(ones((n, n)), dtype=int8) - A
    rows, cols = iA.nonzero()
    count = 0
    for i, j in filter(lambda x: x[0] < x[1], zip(rows, cols)):
        if abs(R[i, j]) <= threshold:   # TODO: find out why negative resistance values
            count += 1
    return count


def matrixToDot(A, file):
    n = A.shape[0]
    file = open("../simulation/" + file, "w")
    print("graph {", file=file, flush=True)
    for i in range(n):
        flag = False
        for j in range(i + 1, n):
            if A[i, j] > 0:
                if not flag: print("\t", end="", file=file)
                flag = True
                print(i, "--", j, ";", sep="", end=" ", file=file)
        if flag: print("", file=file, flush=True)
    print("}", file=file, flush=True)


def degreeDistribution(A):
    degrees = A.sum(1)
    n = degrees.size
    P = zeros(n)
    for i in range(n):
        P[degrees[i]] += 1
    return P / n


def simulate(type, count, p=0.5, start=0, stop=1, step=0.001, mean=1):
    file = open("../simulation/" + str(type) + "_p" + str(p) + "_c" + str(count) + "_m" + str(mean) + ".txt", "w")
    for threshold in arange(start, stop, step):
        sum = 0
        for i in range(mean):
            A = genER(count, p)
            T = sgraph.laplacian(A)
            R = resistance(T.todense())   # TODO: remove weak point
            links = hiddenLinks(A, R, threshold)
            sum += links
        avg = sum / mean
        print(threshold, avg, sep="\t", file=file, flush=True)
        print(threshold, avg, sep="\t")
    file.close()


# start = time.time()
# A = genER(400)
# print(time.time() - start)
# print(A.nnz)

# print(time.ctime(time.time()))
# simulate("ER", 50, 0.4, mean=10)
# print(time.ctime(time.time()))

n = 100
m = int(n * 0.2)
A = genSF(n)
P = degreeDistribution(A)

axes = figure().gca()
axes.set_xticks(arange(0, m + 1, int(m * 0.05)))
axes.set_yticks(arange(0, 1.1, 0.05))
title('Degree Distribution of SF n=' + str(n))
ylabel('P(k)')
xlabel('k')
xlim(0, m)
plot(range(1, m), [P[k] for k in range(1, m)], 'ro')
plot(arange(1, m, 0.01), [pow(k, -3) for k in arange(1, m, 0.01)], 'r-', color='green')
grid()

savefig('../simulation/degree_distribution_SF_' + str(n) + '.png')
show()

#Trying to get determinant in different ways:
# LU = slinalg.splu(T)
# L, U, P, Q, R, do_recip = umf().lu(T)
#
# detT = 1
# for i in range(n):
#     detT *= detT * L[i, i] * U[i, i]
#
# print(detT)
# print(linalg.det(T.todense()))